package org.idea.irpc.framework.core.server;

import io.netty.channel.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collections;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.LongAdder;

/**
 * 最大连接数限制处理器
 *一个基于 Netty 的最大连接数限制处理器，用于控制服务器允许的最大并发连接数，防止因过多连接导致资源耗尽
 * 该处理器通过统计当前活跃连接数，在新连接建立时判断是否超出上限，若超出则拒绝连接并记录日志
 * @Author linhao
 * @Date created in 10:47 上午 2022/3/6
 */
//允许该处理器在多个 ChannelPipeline 中共享。
@ChannelHandler.Sharable
public class MaxConnectionLimitHandler extends ChannelInboundHandlerAdapter {

    private static final Logger LOGGER = LoggerFactory.getLogger(MaxConnectionLimitHandler.class);

    private final int maxConnectionNum;// 最大允许连接数
    private final AtomicInteger numConnection = new AtomicInteger(0);// 当前连接数（原子操作保证线程安全）
    private final Set<Channel> childChannel = Collections.newSetFromMap(new ConcurrentHashMap<>());// 活跃连接的Channel集合
    //这是在jdk1.8之后出现的对于AtomicLong的优化版本
    private final LongAdder numDroppedConnections = new LongAdder();// 被拒绝的连接计数（LongAdder 在竞争激烈时性能更优）
    private final AtomicBoolean loggingScheduled = new AtomicBoolean(false);// 控制日志记录频率的标志

    public MaxConnectionLimitHandler(int maxConnectionNum) {
        this.maxConnectionNum = maxConnectionNum;
    }

    @Override
    public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
        Channel channel = (Channel) msg;
        int conn = numConnection.incrementAndGet();
        if (conn > 0 && conn <= maxConnectionNum) {
            this.childChannel.add(channel);
            channel.closeFuture().addListener(future -> {
                childChannel.remove(channel);
                numConnection.decrementAndGet();
            });
            super.channelRead(ctx, msg);
        } else {
            numConnection.decrementAndGet();

            channel.config().setOption(ChannelOption.SO_LINGER, 0); //立即关闭连接,避免产生大量的time_wait连接
            channel.unsafe().closeForcibly();
            numDroppedConnections.increment();
            //这里加入一道cas可以减少一些并发请求的压力,定期地执行一些日志打印
            if (loggingScheduled.compareAndSet(false, true)) {
                ctx.executor().schedule(this::writeNumDroppedConnectionLog,1, TimeUnit.SECONDS);
            }
        }
    }

    /**
     * 记录连接失败的日志
     * 记录被拒绝的连接数，避免频繁打日志影响性能。
     * 通过 AtomicBoolean 确保只有一个线程在 1 秒内执行日志记录，其他线程跳过
     */
    private void writeNumDroppedConnectionLog() {
        loggingScheduled.set(false);
        final long dropped = numDroppedConnections.sumThenReset();
        if(dropped>0){
            LOGGER.error("Dropped {} connection(s) to protect server,maxConnection is {}",dropped,maxConnectionNum);
        }
    }

}
